from numpy import array,dot,sqrt
import matplotlib
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from random import random

def color_vector(k):
#creates a vector of k different colors
    def f(x): return [(float(x)/float(k)+random())/6.0,(float(x)/float(k)+random())/6.0,(float(x)/float(k)+random())/6.0]
    return map(tuple,map(f,range(k)))


def kmeans_plot2D(points,centers_point,assignments):
    #points is a set of n array with dimension d
    #centers_point is a set of k array of dimension d
    #assignments is a set of n integer with value from 0 to k-1
    n = len(points)
    if(n<2):
        return

    dim = len(points[0])
    if(dim!=2):
        return

    k=len(centers_point)
    colors=color_vector(k)

    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.scatter([0],[0], marker='o')
    for i in range(n):
        ax.plot(points[i][0],points[i][1], color=colors[assignments[i]],  marker='o')
    for i in range(k):
        ax.plot(centers_point[i][0],centers_point[i][1], color=colors[i], marker='x')

    plt.show()

def plotKmeans3D(points,centers_point,assignments):
    #points is a set of n array with dimension d
    #centers_point is a set of k array of dimension d
    #assignments is a set of n integer with value from 0 to k-1
    n = len(points)
    if(n<2):
        return

    dim = len(points[0])
    if(dim!=3):
        return

    k=len(centers_point)
    colors=color_vector(k)

    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')

    ax.scatter([0],[0],[0], marker='o')
    for i in range(n):
        ax.scatter(points[i][0],points[i][1],points[i][2], color=colors[assignments[i]],  marker='o')
    for i in range(k):
        ax.scatter(centers_point[i][0],centers_point[i][1],centers_point[i][2], color=colors[i], marker='x')

    ax.set_xlabel('R')
    ax.set_ylabel('G')
    ax.set_zlabel('B')

    plt.show()
